#!/usr/bin/env bun

import { parse } from "proto-parser"

const typeToDecoder = {
	double: "decodeDouble",
	float: "decodeFloat",
	int32: "decodeInt32",
	uint32: "decodeUint32",
	sint32: "decodeSint32",
	fixed32: "decodeFixed32",
	sfixed32: "decodeSfixed32",
	int64: "decodeInt64",
	uint64: "decodeUint64",
	sint64: "decodeSint64",
	fixed64: "decodeFixed64",
	sfixed64: "decodeSfixed64",
	bool: "decodeBool",
	string: "decodeString",
	bytes: "decodeBytes",
}

const genLiDecoder = (decoderName) => {
	if (decoderName.endsWith("Li")) {
		return decoderName
	}
	return `${decoderName}Li`
}

const Decode = (fieldDecoders, fields) => {
	return (buf, pos = 0) => {
		const result = fields.map((f) => f.defaultValue)
		const oneofSeen = {}

		while (pos < buf.length) {
			const [tag, newPos] = decodeInt32(buf, pos)
			pos = newPos
			const fieldNum = tag >> 3
			const wireType = tag & 7

			const fieldInfo = fieldDecoders[fieldNum]

			if (fieldInfo) {
				let value, nextPos
				if (fieldInfo.isMap) {
					const [mapEntry, entryEnd] = fieldInfo.decoder(buf, pos)
					const map = result[fieldInfo.index] || (result[fieldInfo.index] = {})
					map[mapEntry[0]] = mapEntry[1]
					nextPos = entryEnd
				} else if (fieldInfo.isRepeated) {
					const list = result[fieldInfo.index] || (result[fieldInfo.index] = [])
					if (fieldInfo.isPacked) {
						const [values, listEnd] = fieldInfo.decoder(buf, pos)
						list.push(...values)
						nextPos = listEnd
					} else {
						const [val, itemEnd] = fieldInfo.decoder(buf, pos)
						list.push(val)
						nextPos = itemEnd
					}
				} else {
					;[value, nextPos] = fieldInfo.decoder(buf, pos)
					if (fieldInfo.oneof) {
						if (oneofSeen[fieldInfo.oneof] !== undefined) {
							result[oneofSeen[fieldInfo.oneof]] = undefined
						}
						oneofSeen[fieldInfo.oneof] = fieldInfo.index
					}
					result[fieldInfo.index] = value
				}
				pos = nextPos
			} else {
				pos = skipField(buf, pos, wireType)
			}
		}
		return [result, pos]
	}
}

const generateMessageDecoder = (message, knownTypes, path) => {
	let code = ""
	const messageName = message.name
	const decoderName = `${path.length > 0 ? path.join("_") + "_" : ""}${messageName}Decode`

	const fields = Object.values(message.fields)
	fields.sort((a, b) => a.id - b.id)

	const fieldDecoders = {}
	const fieldArray = []

	for (const field of fields) {
		const fieldInfo = {
			index: field.id - 1,
			oneof: field.oneof || null,
			isRepeated: field.repeated,
			isMap: field.map,
			isPacked:
				field.repeated &&
				[
					"double",
					"float",
					"int32",
					"int64",
					"uint32",
					"uint64",
					"sint32",
					"sint64",
					"fixed32",
					"fixed64",
					"sfixed32",
					"sfixed64",
					"bool",
				].includes(field.type.value),
		}

		let decoder
		if (typeToDecoder[field.type.value]) {
			decoder = typeToDecoder[field.type.value]
			if (field.repeated && !fieldInfo.isPacked) {
				// non-packed use single decoder
			} else if (field.repeated) {
				decoder = genLiDecoder(decoder)
			}
		} else {
			const nestedType = knownTypes[field.type.value]
			if (nestedType) {
				decoder = nestedType.decoderName
			} else {
				console.warn(`Unknown type: ${field.type.value} in ${messageName}`)
				continue
			}
		}

		fieldInfo.decoder = decoder

		if (field.map) {
			const keyType = field.type.key
			const valueType = field.type.value
			const mapEntryName = `${messageName}_${field.name}_Entry`
			const mapEntryDecoderName = `${decoderName}_${field.name}_EntryDecode`

			const mapEntryMessage = {
				name: mapEntryName,
				fields: {
					key: { id: 1, type: { value: keyType } },
					value: { id: 2, type: { value: valueType } },
				},
				nested: {},
			}
			code += generateMessageDecoder(mapEntryMessage, knownTypes, [])
			fieldInfo.decoder = mapEntryDecoderName
		}

		fieldDecoders[field.id] = fieldInfo

		let defaultValue
		if (field.repeated) {
			defaultValue = []
		} else if (field.map) {
			defaultValue = {}
		} else if (
			field.type.value.startsWith("int") ||
			field.type.value.startsWith("uint") ||
			field.type.value.startsWith("sint") ||
			field.type.value.startsWith("fixed")
		) {
			defaultValue = 0
		} else if (field.type.value === "bool") {
			defaultValue = false
		} else if (field.type.value === "string" || field.type.value === "bytes") {
			defaultValue = ""
		} else {
			defaultValue = null
		}
		fieldArray.push({ name: field.name, defaultValue })
	}

	const fieldDecodersString = JSON.stringify(fieldDecoders, null, 2).replace(
		/"decoder": \"([^\"]+)\" /g,
		'"decoder": $1',
	)

	code += `export const ${decoderName} = Decode(${fieldDecodersString}, ${JSON.stringify(
		fieldArray,
		null,
		2,
	)})

`

	if (message.nested) {
		for (const nestedName in message.nested) {
			const nestedMessage = message.nested[nestedName]
			if (nestedMessage.nodeType === "message") {
				const nestedPath = [...path, messageName]
				const { code: nestedCode, knownTypes: updatedKnownTypes } =
					generateMessageDecoder(nestedMessage, knownTypes, nestedPath)
				code += nestedCode
				Object.assign(knownTypes, updatedKnownTypes)
			}
		}
	}

	knownTypes[message.name] = { decoderName }
	if (path.length > 0) {
		knownTypes[`${path.join(".")}.${message.name}`] = { decoderName }
	}

	return { code, knownTypes }
}

const findMessages = (node, path = []) => {
	let messages = {}
	if (node.nodeType === "message") {
		const fqName = [...path, node.name].join(".")
		messages[fqName] = node
		path.push(node.name)
	}

	if (node.nested) {
		for (const key in node.nested) {
			Object.assign(messages, findMessages(node.nested[key], path))
		}
	}
	if (node.nodeType === "message") {
		path.pop()
	}
	return messages
}

export default (proto) => {
	const ast = parse(proto).root

	const imports = new Set(Object.values(typeToDecoder))
	imports.add("skipField")
	imports.add("liDecode")

	let generatedCode = ""
	const allMessages = findMessages(ast)
	const knownTypes = {}

	// Pre-declare all decoder names
	for (const fqName in allMessages) {
		const parts = fqName.split(".")
		const messageName = parts.pop()
		const path = parts
		const decoderName = `${path.length > 0 ? path.join("_") + "_" : ""}${messageName}Decode`
		knownTypes[fqName] = { decoderName }
	}

	// Handle google protobuf special types
	const googleTimestamp = {
		name: "Timestamp",
		fields: {
			seconds: { id: 1, type: { value: "int64" } },
			nanos: { id: 2, type: { value: "int32" } },
		},
		nested: {},
	}
	const { code: tsCode } = generateMessageDecoder(googleTimestamp, knownTypes, [
		"google_protobuf",
	])
	generatedCode += tsCode
	knownTypes["google.protobuf.Timestamp"] = {
		decoderName: "google_protobuf_TimestampDecode",
	}

	const googleAny = {
		name: "Any",
		fields: {
			type_url: { id: 1, type: { value: "string" } },
			value: { id: 2, type: { value: "bytes" } },
		},
		nested: {},
	}
	const { code: anyCode } = generateMessageDecoder(googleAny, knownTypes, [
		"google_protobuf",
	])
	generatedCode += anyCode
	knownTypes["google.protobuf.Any"] = {
		decoderName: "google_protobuf_AnyDecode",
	}

	for (const fqName in allMessages) {
		const message = allMessages[fqName]
		const parts = fqName.split(".")
		parts.pop() // remove message name
		const { code } = generateMessageDecoder(message, knownTypes, parts)
		generatedCode += code
	}

	const finalCode = `
import {
    ${[...new Set([...imports, ...Object.values(typeToDecoder)])].join(",\n\t")}
} from './decode.js'

const Decode = (fieldDecoders, fields) => {
	return (buf, pos = 0) => {
		const result = fields.map(f => {
            if (f.defaultValue && Array.isArray(f.defaultValue)) return []
            if (f.defaultValue && typeof f.defaultValue === 'object') return {}
            return f.defaultValue
        })
		const oneofSeen = {}

		while (pos < buf.length) {
			const [tag, newPos] = decodeInt32(buf, pos)
			pos = newPos
			const fieldNum = tag >> 3
			const wireType = tag & 7

			const fieldInfo = fieldDecoders[fieldNum]

			if (fieldInfo) {
				let value, nextPos
				if (fieldInfo.isMap) {
					const [mapEntry, entryEnd] = fieldInfo.decoder(buf, pos)
					const map = result[fieldInfo.index] || (result[fieldInfo.index] = {})
					map[mapEntry[0]] = mapEntry[1]
					nextPos = entryEnd
				} else if (fieldInfo.isRepeated) {
                    const list = result[fieldInfo.index]
					if (fieldInfo.isPacked) {
						const [values, listEnd] = fieldInfo.decoder(buf, pos)
						list.push(...values)
						nextPos = listEnd
					} else {
						const [val, itemEnd] = fieldInfo.decoder(buf, pos)
						list.push(val)
						nextPos = itemEnd
					}
				} else {
					;[value, nextPos] = fieldInfo.decoder(buf, pos)
					if (fieldInfo.oneof) {
						if (oneofSeen[fieldInfo.oneof] !== undefined && oneofSeen[fieldInfo.oneof] !== fieldInfo.index) {
							result[oneofSeen[fieldInfo.oneof]] = fields[oneofSeen[fieldInfo.oneof]].defaultValue
						}
						oneofSeen[fieldInfo.oneof] = fieldInfo.index
					}
					result[fieldInfo.index] = value
				}
				pos = nextPos
			} else {
				pos = skipField(buf, pos, wireType)
			}
		}
		return [result, pos]
	}
}

${generatedCode}
`

	return finalCode
}
